#include "FluidSimulator.h"
#include <malloc.h>
#include <math.h>

#pragma comment(lib,"RFFTW2dll.lib")

FluidSimulator::FluidSimulator(void)
{
}

FluidSimulator::~FluidSimulator(void)
{
}

void FluidSimulator::InitSimulator( int size, fftw_real visc )
{
	scale = size;
	this->visc = visc;

	u = (fftw_real*)malloc(scale*(scale+2)*sizeof(fftw_real));
	v = (fftw_real*)malloc(scale*(scale+2)*sizeof(fftw_real));
	u_old = (fftw_real*)malloc(scale*(scale+2)*sizeof(fftw_real));
	v_old = (fftw_real*)malloc(scale*(scale+2)*sizeof(fftw_real));

	d = (fftw_real*)malloc(scale*scale*sizeof(fftw_real));
	d_old = (fftw_real*)malloc(scale*scale*sizeof(fftw_real));

	f2u = (fftw_real*)malloc(scale*scale*sizeof(fftw_real));
	f2v = (fftw_real*)malloc(scale*scale*sizeof(fftw_real));

	for (int i=0;i<scale*(scale+2);i++)
	{
		u[i] = v[i] = u_old[i] = v_old[i] = 0;
	}

	for (int i=0;i<scale*scale;i++)
	{
		d[i] = d_old[i] = f2u[i] = f2v[i] = 0;
	}

	InitFFT();

}

void FluidSimulator::InitFFT()
{
	r2c = rfftw2d_create_plan(scale, scale, FFTW_REAL_TO_COMPLEX, FFTW_IN_PLACE);
	c2r = rfftw2d_create_plan(scale, scale, FFTW_COMPLEX_TO_REAL, FFTW_IN_PLACE);
}

void FluidSimulator::r2cFFT( fftw_real** u )
{
	rfftwnd_one_real_to_complex(r2c,(fftw_real *)*u,(fftw_complex*)*u);
}

void FluidSimulator::c2rFFT( fftw_real** u )
{
	rfftwnd_one_complex_to_real(c2r,(fftw_complex *)*u,(fftw_real *)*u);
}

void FluidSimulator::DiffussionProjection()
{
	fftw_real x, y, f, r, U[2], V[2];
	int i, j, n;

	n = scale;

	for ( i=0 ; i<=n ; i+=2 ) {
		x = 0.5f*i;
		for ( j=0 ; j<n ; j++ ) {
			y = j<=n/2 ? (fftw_real)j : (fftw_real)j-n;
			r = x*x+y*y;
			if ( r==0.0f ) continue;
			f = (fftw_real)exp(-r*dt*visc);
			U[0] = u_old[i  +(n+2)*j]; V[0] = v_old[i  +(n+2)*j];
			U[1] = u_old[i+1+(n+2)*j]; V[1] = v_old[i+1+(n+2)*j];

			u_old[i  +(n+2)*j] = f*( (1-x*x/r)*U[0]     -x*y/r *V[0] );
			u_old[i+1+(n+2)*j] = f*( (1-x*x/r)*U[1]     -x*y/r *V[1] );
			v_old[i+  (n+2)*j] = f*(   -y*x/r *U[0] + (1-y*y/r)*V[0] );
			v_old[i+1+(n+2)*j] = f*(   -y*x/r *U[1] + (1-y*y/r)*V[1] );
		}    
	}
}

void FluidSimulator::VelocityAdvection()
{
	fftw_real x, y, x0, y0, s, t;
	int i, j, i0, j0, i1, j1, n;

	n = scale;

	for ( x=0.5f/n,i=0 ; i<n ; i++,x+=1.0f/n ) 
	{
		for ( y=0.5f/n,j=0 ; j<n ; j++,y+=1.0f/n ) 
		{
			x0 = n*(x-dt*u_old[i+n*j])-0.5f; 
			y0 = n*(y-dt*v_old[i+n*j])-0.5f;
			i0 = floor(x0);
			s = x0-i0;
			i0 = (n+(i0%n))%n;
			i1 = (i0+1)%n;
			j0 = floor(y0);
			t = y0-j0;
			j0 = (n+(j0%n))%n;
			j1 = (j0+1)%n;
			u[i+n*j] = (1-s)*((1-t)*u_old[i0+n*j0]+t*u_old[i0+n*j1])+                        
				s *((1-t)*u_old[i1+n*j0]+t*u_old[i1+n*j1]);
			v[i+n*j] = (1-s)*((1-t)*v_old[i0+n*j0]+t*v_old[i0+n*j1])+
				s *((1-t)*v_old[i1+n*j0]+t*v_old[i1+n*j1]);
		}    
	} 
}

void FluidSimulator::VelocitySwap()
{
	int n;

	n = scale;

	for ( int i=0 ; i<n ; i++ )
	{
		for ( int j=0 ; j<n ; j++ ) { 
			u_old[i+(n+2)*j] = u[i+n*j]; 
			v_old[i+(n+2)*j] = v[i+n*j];
		}
	}
}

void FluidSimulator::VelocityNormalize()
{
	int n;
	fftw_real f;
	
	n = scale;

	f = 1.0/(n*n);
	for ( int i=0 ; i<n ; i++ )
	{
		for ( int j=0 ; j<n ; j++ ) {
			u[i+n*j] = f*u_old[i+(n+2)*j]; 
			v[i+n*j] = f*v_old[i+(n+2)*j]; 
		}
	}
		
}

void FluidSimulator::DensityAdvection()
{
	fftw_real x, y, x0, y0, s, t;
	int i, j, i0, j0, i1, j1, n;

	n = scale;

	for ( x=0.5f/n,i=0 ; i<n ; i++,x+=1.0f/n ) 
	{
		for ( y=0.5f/n,j=0 ; j<n ; j++,y+=1.0f/n ) 
		{
			x0 = n*(x-dt*u[i+n*j])-0.5f; 
			y0 = n*(y-dt*v[i+n*j])-0.5f;
			i0 = floor(x0);
			s = x0-i0;
			i0 = (n+(i0%n))%n;
			i1 = (i0+1)%n;
			j0 = floor(y0);
			t = y0-j0;
			j0 = (n+(j0%n))%n;
			j1 = (j0+1)%n;
			d[i+n*j] = (1-s)*((1-t)*d_old[i0+n*j0]+t*d_old[i0+n*j1])+                        
				s *((1-t)*d_old[i1+n*j0]+t*d_old[i1+n*j1]);
		}    
	} 
}

void FluidSimulator::AddExternalForce()
{
	for (int i = 0; i < scale * scale; i++) {
		d_old[i] = d[i];

		f2u[i] *= 0.85;
		f2v[i] *= 0.85;

		u_old[i] = f2u[i];
		v_old[i] = f2v[i];

		u[i] += dt*u_old[i]; 
		u_old[i] = u[i];

		v[i] += dt*v_old[i]; 
		v_old[i] = v[i];
	}  
}

void FluidSimulator::Update( fftw_real timestep )
{
	dt = timestep;

	UpdateVelocity();

	UpdateDensity();	
}

void FluidSimulator::UpdateVelocity()
{
	AddExternalForce();

	VelocityAdvection();

	VelocitySwap();

	r2cFFT(&u_old);
	r2cFFT(&v_old);

	DiffussionProjection();

	c2rFFT(&u_old);
	c2rFFT(&v_old);

	VelocityNormalize();
}

void FluidSimulator::UpdateDensity()
{
	DensityAdvection();
}
